package com.xzc.kafka.action.producer;

import org.apache.kafka.clients.producer.*;
import org.apache.kafka.common.Cluster;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.serialization.Serializer;
import org.apache.kafka.common.serialization.StringSerializer;

import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Random;

/**
 * 将包含有audit key的消息发送到最后一个分区
 */
public class _01AuditPartitioner implements Partitioner {
    private Random random;

    @Override
    public int partition(String topic, Object keyObj, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        String key = (String) keyObj;
        List<PartitionInfo> partitionInfoList = cluster.availablePartitionsForTopic(topic);
        int partitionCount = partitionInfoList.size();
        int auditPartition = partitionCount - 1;
        return key == null || key.isEmpty() || !key.contains("audit")
                ? random.nextInt(partitionCount - 1) : auditPartition;
    }

    @Override
    public void close() {

    }

    @Override
    public void configure(Map<String, ?> configs) {
        random = new Random();
    }

    public static void main(String[] args) throws Exception {
        Properties props = new Properties();
        props.put("bootstrap.servers", "es1:19092,es2:19092,es3:19092");
//        props.put("key.serializer", "org.apache.kafka.common.serializtion.StringSerializer");
//        props.put("value.serializer", "org.apache.kafka.common.serializtion.StringSerializer");
        props.put("acks", "-1");
        props.put("retries", 3);
        props.put("batch.size", 323840);
        props.put("linger.ms", 10);
        props.put("buffer.memory", 33554432);
        props.put("max.block.ms", 3000);

        props.put(ProducerConfig.PARTITIONER_CLASS_CONFIG, "com.xzc.kafka.action.producer._01AuditPartitioner");

        Serializer<String> keySerl = new StringSerializer();
        Serializer<String> valSerl = new StringSerializer();
        Producer<String, String> producer = new KafkaProducer<>(props, keySerl, valSerl);

        String topic = "test-topic1";
        ProducerRecord nonKeyRecord = new ProducerRecord(topic, "non-key record");
        ProducerRecord auditRecord = new ProducerRecord(topic, "audit", "audit record");
        ProducerRecord nonAuditRecord = new ProducerRecord(topic, "other", "non-audit record");
        producer.send(nonKeyRecord).get();
        producer.send(nonAuditRecord).get();
        producer.send(auditRecord).get();
        producer.send(nonKeyRecord).get();
        producer.send(nonAuditRecord).get();

        producer.close();
    }
}
